'''
Created on Feb 28, 2013

@author: swLi, zhaonan
'''
import numpy as np

from MatViewer.MatViewer import MatViewer

from traits.api import HasTraits, Instance, Float
from phy.general.Constant import GammaUnit_B, Gamma_C13
from phy.general.MaterialData import NV_AXIS

from phy.solver.eigen import Eigen_NSpinHalf
from phy.interaction.SingleSpinInteraction import SingleSpinInteraction
from phy.interaction.MultiSpinInteraction import  MultiSpinInteraction
from phy.spin.Lattice import DiamondLattice

from math_qm.vector import vector
from math_qm.Constant import X, Y, Z

class NV_Bath_System(HasTraits):
    B0   = Instance(vector)
    B_NV = Instance(vector)
    Omega_Rabi = Float()
    NV_detune  = Float()
        
    bath_abundance = Float()
    bath_Rmax      = Float()
    bath_seed      = Float()
    bath_chop_dis  = Float()    
    
    NV      = Instance(SingleSpinInteraction)
    bath    = Instance(DiamondLattice)
    nv_bath = Instance(MultiSpinInteraction)
    
    def __init__(self,applied_B, Omega_Rabi, NV_detune, abundance, Rmax, chop_dist, seednum):
        self.B0   = vector(xyz= [        0., 0., Gamma_C13 * GammaUnit_B * Applied_B])
        self.B_NV = vector(xyz= [Omega_Rabi, 0., NV_detune])
        
        self.Omega_Rabi = Omega_Rabi
        self.NV_detune  = NV_detune
        
        self.bath_abundance = abundance
        self.bath_Rmax      = Rmax
        self.bath_chop_dist = chop_dist
        self.bath_seed      = seednum
        
        self.generate_system()
    
    def generate_NV(self):
        self.NV = SingleSpinInteraction(element='E', spin_type='Central', spin_mark='None', position=[0., 0., 0.])
        print "NV center is generated."
    
    def generate_bath(self):
        self.bath = DiamondLattice()
        self.bath.impurity(abundance=self.bath_abundance, Rmax=self.bath_Rmax, seednum = self.bath_seed)
        self.bath.chop(self.bath_chop_dist)

    
    def set_bath_local_field(self):
        print "calculating local field of bath spins."
        for i in range(1, self.bath.impurity_num+1):
            dip     = self.nv_bath.dipolar_tensor(0, i, axes=NV_AXIS)
            local_field_value = np.array([dip.zx, dip.zy, dip.zz]) + self.B0.xyz
            local_b = vector( xyz =  local_field_value)
            
            self.nv_bath.spin_list[i].set_local_field(local_b)    
    
    def generate_system(self):
        self.generate_NV()
        
        self.generate_bath()
        
        print "generating composed system."
        spin_list = [self.NV]
        for pos in self.bath.impurity_pos:
            c13spin = SingleSpinInteraction(element='C13', spin_type='Bath', spin_mark='None', position=pos)
            spin_list.append(c13spin)
        self.nv_bath = MultiSpinInteraction(spins = spin_list)
        
        self.NV.set_local_field(B_NV)
        self.set_bath_local_field()

class Eigen_NVbath(Eigen_NSpinHalf):
    """
    To generate a driving NV interacting with many impurities
    """
    
    def __init__(self, spins, is_auto_update=False):
        Eigen_NSpinHalf.__init__(self, spins, is_auto_update=is_auto_update)
        
    def scattering_amplitude(self, idx, dim, states):
        """
        This method is originally defined in father class. Here it should be replaced.
        """
        len_states = len(idx)
        
        if len_states == 2 and 0 in idx:
            i, j = idx
            dip = self.spins.dipolar_tensor(i, j, axes= NV_AXIS)
            res = 0.25 * ( dip.zx * np.kron(Z,X) \
                         + dip.zy * np.kron(Z,Y) \
                         + dip.zz * np.kron(Z,Z) )
            """ res = S_i * Dipo_ij * S_j. Here S_0=(0,0,Sz) """
        else:
            res = Eigen_NSpinHalf.scattering_amplitude(self, idx, dim, states)
        
        return res


if __name__ == '__main__':
  
    Applied_B  = 600.0
    Omega_Rabi = 600.0
    NV_detune  = 10.0
    
    abundance = 0.011
    Rmax      = 4
    dist      = 10.0
    seednum   = 1
    
    B0   = vector(xyz= [        0., 0., Applied_B])
    B_NV = vector(xyz= [Omega_Rabi, 0., NV_detune])
    
    nv_bath_sys=NV_Bath_System(applied_B = Applied_B,\
                               Omega_Rabi= Omega_Rabi,\
                               NV_detune = NV_detune, \
                               abundance = abundance, \
                               Rmax      = Rmax, \
                               chop_dist = dist, \
                               seednum   = seednum)
    
    print "%d bath spins are generated!" % nv_bath_sys.bath.impurity_num
    nv_plus_bath = nv_bath_sys.nv_bath

    spin_eigs = Eigen_NVbath(spins=nv_plus_bath)
    print "calculating Hamiltonian matrix..."
    spin_eigs.calculate_mat()
    print spin_eigs.mat
    
    mv = MatViewer(spin_eigs.mat)
    mv.configure_traits()
